#include <unittest/unittest.h>

#include <cusp/io/matrix_market.h>
#include <cusp/gallery/poisson.h>
#include <cusp/multiply.h>
#include <cusp/krylov/gmres.h>
#include <cusp/precond/diagonal.h>

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <sys/time.h>

template <class MemorySpace>
void TestGmresMatrixMarket(void)
{
	typedef double ValueType;
	
    cusp::csr_matrix<int, ValueType, MemorySpace> A;
    cusp::io::read_matrix_market_file(A, "/project/csstaff/inputs/Matrices/MatrixMarket/SMALL/jpwh_991.mtx");

    cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 0.0);
    cusp::array1d<ValueType, MemorySpace> b(A.num_rows, 1.0);

	cusp::default_monitor<ValueType> monitor(b, 50, 0.0, 1e-6);

    cusp::krylov::gmres(A, x, b, 50, monitor);

    // check residual norm
    cusp::array1d<ValueType, MemorySpace> residual(A.num_rows, 0.0);
    cusp::multiply(A, x, residual);
    cusp::blas::axpby(residual, b, residual, -1.0, 1.0);

    ASSERT_EQUAL(cusp::blas::nrm2(residual) < 1e-6 * cusp::blas::nrm2(b), true);
	ASSERT_EQUAL(monitor.iteration_count() < 45, true);
}
DECLARE_HOST_DEVICE_UNITTEST(TestGmresMatrixMarket);

template <class MemorySpace>
void TestGmresRestart(void)
{
	typedef double ValueType;
	
    cusp::csr_matrix<int, ValueType, MemorySpace> A;
    cusp::io::read_matrix_market_file(A, "/project/csstaff/inputs/Matrices/MatrixMarket/SMALL/jpwh_991.mtx");

    cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 0.0);
    cusp::array1d<ValueType, MemorySpace> b(A.num_rows, 1.0);

	cusp::default_monitor<ValueType> monitor(b, 50, 0.0, 1e-6);

    cusp::krylov::gmres(A, x, b, 40, monitor);

    // check residual norm
    cusp::array1d<ValueType, MemorySpace> residual(A.num_rows, 0.0);
    cusp::multiply(A, x, residual);
    cusp::blas::axpby(residual, b, residual, -1.0, 1.0);

    ASSERT_EQUAL(cusp::blas::nrm2(residual) < 1e-6 * cusp::blas::nrm2(b), true);
	ASSERT_EQUAL(monitor.iteration_count() < 50, true);
}
DECLARE_HOST_DEVICE_UNITTEST(TestGmresRestart);

template <class MemorySpace>
void TestGmresPoisson(void)
{
	typedef double ValueType;
	
    cusp::csr_matrix<int, ValueType, MemorySpace> A;

    cusp::gallery::poisson5pt(A, 5, 5);

    cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 0.0);
    cusp::array1d<ValueType, MemorySpace> b(A.num_rows, 1.0);

    cusp::default_monitor<ValueType> monitor(b, 20, 0.0, 1e-4);
    
    cusp::krylov::gmres(A, x, b, 20, monitor);

    // check residual norm
    cusp::array1d<ValueType, MemorySpace> residual(A.num_rows, 0.0);
    cusp::multiply(A, x, residual);
    cusp::blas::axpby(residual, b, residual, -1.0, 1.0);

    ASSERT_EQUAL(cusp::blas::nrm2(residual) < 1e-4 * cusp::blas::nrm2(b), true);
}
DECLARE_HOST_DEVICE_UNITTEST(TestGmresPoisson);

void TestGmresZeroResidualHost(void)
{
	typedef typename cusp::host_memory MemorySpace;
	typedef double ValueType;
	
    cusp::array2d<ValueType, MemorySpace> A(2,2);
    A(0,0) = 8; A(0,1) = 0;
    A(1,0) = 0; A(1,1) = 4;

    cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 1.0);
    cusp::array1d<ValueType, MemorySpace> b(A.num_rows);

    cusp::multiply(A, x, b);

    cusp::default_monitor<ValueType> monitor(b, 20, 0.0);
    
    cusp::krylov::gmres(A, x, b, 20, monitor);

    // check residual norm
    cusp::array1d<ValueType , MemorySpace> residual(A.num_rows, 0.0);
    cusp::multiply(A, x, residual);
    cusp::blas::axpby(residual, b, residual, -1.0, 1.0);

    ASSERT_EQUAL(monitor.converged(),        true);
    ASSERT_EQUAL(monitor.iteration_count(),     0);
    ASSERT_EQUAL(cusp::blas::nrm2(residual), 0.0);
}
DECLARE_UNITTEST(TestGmresZeroResidualHost);

void TestGmresZeroResidualDevice(void)
{
	typedef typename cusp::device_memory MemorySpace;
	typedef double ValueType;
	
    cusp::array2d<ValueType, MemorySpace, cusp::column_major> dA(2,2);
    dA(0,0) = 8; dA(0,1) = 0;
    dA(1,0) = 0; dA(1,1) = 4;

    cusp::array1d<ValueType, MemorySpace> dx(dA.num_rows, 1.0);
    cusp::array1d<ValueType, MemorySpace> db(dA.num_rows);
	db[0] = 8; db[1] = 4;

    cusp::default_monitor<ValueType> monitor(db, 20, 0.0);
    
    cusp::krylov::gmres(dA, dx, db, 20, monitor);

	cusp::array2d<ValueType, cusp::host_memory> hA = dA;
	cusp::array1d<ValueType, cusp::host_memory> hx = dx;
	cusp::array1d<ValueType, cusp::host_memory> hb = db;

    // check residual norm
    cusp::array1d<ValueType, cusp::host_memory> residual(hA.num_rows, 0.0);
    cusp::multiply(hA, hx, residual);
    cusp::blas::axpby(residual, hb, residual, -1.0, 1.0);

    ASSERT_EQUAL(monitor.converged(),        true);
    ASSERT_EQUAL(monitor.iteration_count(),     0);
    ASSERT_EQUAL(cusp::blas::nrm2(residual), 0.0);
}
DECLARE_UNITTEST(TestGmresZeroResidualDevice);

template <class MemorySpace>
void TestGmresIdentity(void)
{
	typedef double ValueType;

	cusp::identity_operator<ValueType, MemorySpace> A(10, 10);
    cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 0.0);
    cusp::array1d<ValueType, MemorySpace> b(A.num_rows, 0.0);

    cusp::default_monitor<ValueType> monitor(b, 10, 0);
    
    cusp::krylov::gmres(A, x, b, 10, monitor);

    // check residual norm
    cusp::array1d<ValueType , MemorySpace> residual(A.num_rows, 0.0);
    cusp::multiply(A, x, residual);
    cusp::blas::axpby(residual, b, residual, -1.0, 1.0);

    ASSERT_EQUAL(monitor.converged(), true);
    ASSERT_EQUAL(monitor.iteration_count() <= 10, true);
    ASSERT_EQUAL(cusp::blas::nrm2(residual) <= 0, true);
}
DECLARE_HOST_DEVICE_UNITTEST(TestGmresIdentity);

void TestGmresLaplaceHost(void)
{
	typedef typename cusp::host_memory MemorySpace;
	typedef double ValueType;
	
	const size_t size = 50;

	cusp::array2d<ValueType, MemorySpace> A;
    cusp::gallery::poisson9pt(A, size, size);

	cusp::blas::scal(A.values, 1e-6);
	for(size_t i = 0; i < A.num_rows; i++)
		A(i,i) = A(i,i) + 1;
	
	cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 0.0);
    cusp::array1d<ValueType, MemorySpace> b(A.num_rows, 1.0);

	cusp::default_monitor<ValueType> monitor(b, 100, 0.0, 1e-4);

    cusp::krylov::gmres(A, x, b, 100, monitor);

	// check residual norm
    cusp::array1d<ValueType , MemorySpace> residual(A.num_rows, 0.0);
    cusp::multiply(A, x, residual);
    cusp::blas::axpby(residual, b, residual, -1.0, 1.0);

	ASSERT_EQUAL(cusp::blas::nrm2(residual) < 1e-4 * cusp::blas::nrm2(b), true);
}
DECLARE_UNITTEST(TestGmresLaplaceHost);

void TestGmresLaplaceDevice(void)
{
	typedef typename cusp::device_memory MemorySpace;
	typedef double ValueType;
	
	const size_t size = 50;

	cusp::array2d<ValueType, MemorySpace> dA;
    cusp::gallery::poisson9pt(dA, size, size);

	cusp::blas::scal(dA.values, 1e-6);
	for(size_t i = 0; i < dA.num_rows; i++)
		dA(i,i) = dA(i,i) + 1;
	
	cusp::array1d<ValueType, MemorySpace> dx(dA.num_rows, 0.0);
    cusp::array1d<ValueType, MemorySpace> db(dA.num_rows, 1.0);

	cusp::default_monitor<ValueType> monitor(db, 100, 0.0, 1e-4);

    cusp::krylov::gmres(dA, dx, db, 100, monitor);
	
	cusp::array2d<ValueType, cusp::host_memory> hA = dA;
	cusp::array1d<ValueType, cusp::host_memory> hx = dx;
	cusp::array1d<ValueType, cusp::host_memory> hb = db;

	// check residual norm
    cusp::array1d<ValueType , cusp::host_memory> residual(hA.num_rows, 0.0);
    cusp::multiply(hA, hx, residual);
    cusp::blas::axpby(residual, hb, residual, -1.0, 1.0);
	
	ASSERT_EQUAL(cusp::blas::nrm2(residual) < 1e-4 * cusp::blas::nrm2(hb), true);
}
DECLARE_UNITTEST(TestGmresLaplaceDevice);

